import torch
from FVD.scripts.calc_metrics_for_dataset import calc_metrics_for_dataset
from torch.multiprocessing import Queue
import pandas as pd

# 请使用单GPU计算, 否则好像会出现错误
# 每个视频至少要有25帧. 如果不是这样, 可以去FVD/training/dataset.py修改和25相关的相关代码
# 计算时是从视频所有帧均匀采样25帧, 然后随机offset取16帧进行计算
def calc(real_path, fake_path, queue, gpus=1, mu_real=None, sigma_real=None):
  calc_metrics_for_dataset(None,metrics=['fvd2048_16f'],real_data_path=real_path,
                         fake_data_path=fake_path,mirror=1,resolution=256,
                         gpus=gpus,verbose=0,use_cache=0,num_runs=1,queue=queue,
                         mu_real=mu_real, sigma_real=sigma_real)

if __name__ == '__main__':
  torch.multiprocessing.set_start_method('spawn')
  # 原数据集, 要求目录下全是视频
  input_dir = '11K_output'  
  # 生成数据集, 要求目录下全是视频
  fake_video_dirs = [
        "I2VEdit_output", 
        "AnyV2V_output", 
        "viewcrafter_output", 
        "viewExtra_output",
        "ours",
    ]
  output = 'exp1_fvd.csv'
  fvd_results = []
  
  #用来暂存原数据集的计算结果, 否则每次都要重新计算
  mu_real = None
  sigma_real = None
  
  
  for fake_path in fake_video_dirs:
    print('begin handle:',fake_path)
    que = Queue()
    calc(real_path=input_dir,fake_path=fake_path,queue=que, 
         mu_real=mu_real, sigma_real=sigma_real)
    result = que.get()
    print(f'{fake_path}:',result)
    fvd_results.append([result['results']['fvd2048_16f']])
    mu_real = result['results']['mu_real']
    sigma_real = result['results']['sigma_real']
  
  results_dict = dict(zip(fake_video_dirs, fvd_results))
  
  df = pd.DataFrame(results_dict)
  # os.makedirs(os.path.dirname(output), exist_ok=True)
  df.to_csv(output, index=False)
    